package com.han.agent.controller;

import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.Map;

/**
 * @author han
 */
@RestController
@RequestMapping("/embedding")
public class EmbeddingController {

    @Resource
    private EmbeddingModel embeddingModel;
    @Resource
    private VectorStore vectorStore;

    @GetMapping("/embedding")
    public float[] embedding(@RequestParam(required = false, defaultValue = "Are You OK?") String message) {
        return embeddingModel.embed(message);
    }

    @GetMapping("/store")
    public void store() {
        List<Document> documents = List.of(
                new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!", Map.of("meta1", "meta1")),
                new Document("The World is Big and Salvation Lurks Around the Corner"),
                new Document("You walk forward facing the past and you turn back toward the future.", Map.of("meta2", "meta2")));

        vectorStore.add(documents);
    }

    @GetMapping("/search")
    public List<Document> search() {
        return this.vectorStore.similaritySearch(SearchRequest.builder().query("Spring").topK(5).build());
    }

}
